/*
 * socket.cpp
 *
 *  Created on: 11-May-2011
 *      Author: tan-ce
 */

#include "socket.h"
#include "fcntl.h"
#include <cstdio>

#include <iostream>

using namespace tcetools;

////////////////////////////////////////////////////////////////////////////////

void Socket::chunk::reset(size_t newsize) {
	if (buf) {
		delete [] buf;
	}
	buf = new char[newsize];
	size = newsize;
	offset = 0;
}

////////////////////////////////////////////////////////////////////////////////

void Socket::init() {
	line_delimiter = new char[3];
	line_delimiter[0] = '\r';
	line_delimiter[1] = '\n';
	line_delimiter[2] = '\0';
	line_delimiter_len = 2;

	input_buf_collected = 0;
	socket_closed = false;
}

void Socket::set_noblock() {
	int cur_flags = fcntl(sck_id, F_GETFL, 0);

	if (fcntl(sck_id, F_SETFL, cur_flags | O_NONBLOCK) == -1) {
		SocketError se("Unable to switch to non-blocking mode", this);
		se.append_unix_error();
		throw se;
	}
}

ssize_t Socket::raw_read(char *buf, size_t maxlen) {
	ssize_t ret;
	ret = read(sck_id, (void *) buf, maxlen);
	if (ret < 0) { // Error
		SocketError se("Error reading from Socket!", this);
		se.append_unix_error();
		throw se;
	}

	return ret;
}

ssize_t Socket::raw_write(const char *buf, size_t len) {
	ssize_t ret;
	ret = write(sck_id, (void *) buf, len);
	if (ret < 0) { // Error
		SocketError se("Error writing to Socket!", this);
		se.append_unix_error();
		throw se;
	}

	return ret;
}

void Socket::blocking_unbuffered_read(char *buf, size_t len) {
	ssize_t readlen = 0;
	ssize_t ret;
	
	while (readlen < len) {
		try {
			ret = raw_read(&buf[readlen], len - readlen);
			if (ret == 0) throw SocketError("Socket was closed!", this);
			readlen += ret;
		} catch (SocketError se) {
			if (!se.would_block()) throw se;
		}

		if (readlen >= len) return;

		usleep(1000); // Sleep for 10ms
	}
}

void Socket::buffered_read() {
	ssize_t buffer_left = __SOCKET_BUFFER_SIZE - input_buf_collected;

	if (buffer_left <= 0) {
		SocketError se("Input buffer exceeded", this);
		throw se;
	}

	ssize_t ret = raw_read(&input_buf[input_buf_collected], buffer_left);
	if (ret == 0) {
		SocketError se("Socket closed", this);
		se.socket_closed = true;
		socket_closed = true;
		throw se;
	}

	input_buf_collected += ret;
}

bool Socket::can_read() {
	if (socket_closed) {
		if (input_buf_collected == 0) {
			SocketError se("Socket closed", this);
			se.socket_closed = true;
			throw se;
		} else return true;
	}

	try {
		buffered_read();
	} catch (SocketError se) {
		if (se.would_block()) {
			return (bool) input_buf_collected;
		} else if (socket_closed && input_buf_collected > 0) {
			return true;
		} else throw se;
	}

	return true;
}

void Socket::blocking_write(const char *buf, size_t len) {
	ssize_t ret;
	ssize_t written = 0;

	while (written < len) {
		try {
			ret = raw_write(&buf[written], len - written);
			written += ret;
		} catch (SocketError se) {
			if (!se.would_block()) throw se;
		}

		if (written >= len) return;

		usleep(1000); // Sleep for 10ms
	}
}

void Socket::blocking_read(char *buf, size_t len) {
	while(input_buf_collected < len) {
		try {
			buffered_read();
		} catch (SocketError se) {
			if (se.would_block()) {
				usleep(1000);
				continue;
			} else throw se;
		}
	}

	for (ssize_t i = 0; i < len; i++) buf[i] = input_buf[i];
	shift_input_buf(len);

}

void Socket::http_chunk_write(const char *buf, size_t len) {
	char preamble[7];
	size_t offset = 0;
	
	while (offset < len) {
		int chunk_len = len > 0xffff ? 0xffff : len;
		len -= chunk_len;
		
		snprintf(preamble, 7, "%x\r\n", chunk_len);
		blocking_write(preamble);
		blocking_write(buf + offset, chunk_len);
		
		offset += chunk_len;
	}
}

unsigned int Socket::http_chunk_receive() {
	char preamble[7];
	size_t i = 0, read_len;
	chunk *newchunk;
	
	// Read till \r\n
	do {
		read_len = raw_read(preamble + i, 1);
		if (read_len == 0) {
			throw SocketError("Socket was closed!", this);
		}
		
		i++;
		
		if (i >= 7) {
			throw SocketError(
				"http_chunk_read: Data stream corruption.", this
			);
		}
		
	} while (i < 2 || preamble[i - 2] != '\r' || preamble[i - 1] != '\n');
	preamble[i] = 0;
	
	// Read off chunk size
	sscanf(preamble, "%x\r\n", &chunk_buf.size);
	if (chunk_buf.size < 1) {
		throw SocketError(
			"http_chunk_read: Data stream corruption (invalid chunk size)", this
		);
	} else chunk_buf.reset();
	
	// Now, actually get the chunk
	blocking_unbuffered_read(chunk_buf.buf, chunk_buf.size);
	
	/* Memory debugging:
	for (int i = 0; i < chunk_buf.size; i++) {
		printf("%d\n", i);
		chunk_buf.buf[i] = i;
	} */
	
	return chunk_buf.size;
}

void Socket::http_chunk_read(char *buf, size_t len) {
	unsigned int offset = 0;
	unsigned int len_left = len;
	
	while (offset < len) {
		if (chunk_buf.offset >= chunk_buf.size) http_chunk_receive();
		unsigned int chunk_left = chunk_buf.size - chunk_buf.offset;
		unsigned int tocopy = chunk_left < len_left ? chunk_left : len_left;
		for (int i = 0; i < tocopy; i++) {
			buf[i + offset] = chunk_buf.buf[i + chunk_buf.offset];
		}
		
		offset += tocopy;
		len_left -= tocopy;
		chunk_buf.offset += tocopy;
	}
}

void Socket::set_line_delimiter(const char *new_delim) {
	int len;
	for (len = 0; new_delim[len]; len++);
	delete [] line_delimiter;
	line_delimiter = new char[len + 1];
	line_delimiter_len = len;
	for (int i = 0; i <= len; i++) line_delimiter[i] = new_delim[i];
}

std::string Socket::read_line(bool no_block) {
	bool same;
	std::string ret;
	ssize_t last_unchecked = 0;

	while(1) {
		if (input_buf_collected >= line_delimiter_len) {
			// Look for delimiter
			for (	ssize_t i = last_unchecked;
					i < (input_buf_collected - line_delimiter_len + 1);
					i++) {

				// Check if the delimiter is here
				same = true;
				for (ssize_t delim_pos = 0; delim_pos < line_delimiter_len; delim_pos++) {
					if (line_delimiter[delim_pos] != input_buf[i + delim_pos]) {
						same = false;
						break;
					}
				}

				if (same) { // We've found a line break
					input_buf[i] = '\0';
					ret = (const char *) input_buf;

					// Shift the buffer

					i += line_delimiter_len;

					shift_input_buf(i);

					//std::cout << "Uncollected: " << input_buf_collected << std::endl;

					return ret;
				}
			}
		}

		// Socket closed, and no delimiter in sight. Return what we have.
		if (socket_closed) {
			if (input_buf_collected > 0) {
				input_buf[input_buf_collected] = 0;
				ret = (const char *) input_buf;
				shift_input_buf(input_buf_collected);

				return ret;
			} else {
				SocketError se("Socket closed", this);
				se.socket_closed = true;
				throw se;
			}
		}

		last_unchecked = input_buf_collected - line_delimiter_len;

		bool blocking;
		do {
			ssize_t prev = input_buf_collected;
			can_read();
			blocking = (input_buf_collected == prev);

			if (no_block && blocking) {
				SocketError se("Operation would block", this);
				se.blocked = true;
				throw se;
			}

			usleep(1000);
		} while (blocking);
	}
}

////////////////////////////////////////////////////////////////////////////////

ClientSocket::ClientSocket(const char *rem_host, const char *rem_service) {
	struct hostent	*phe;	/* pointer to host information entry	*/
	struct servent	*pse;	/* pointer to service information entry	*/
	struct protoent *ppe;	/* pointer to protocol information entry*/
	struct sockaddr_in sin;	/* an Internet endpoint address		    */
	int	type;	/* socket descriptor and socket type	*/

	local_ip = "0.0.0.0";

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;

    /* Map service name to port number */
	if ( pse = getservbyname(rem_service, "tcp") )
		sin.sin_port = pse->s_port;
	else if ((sin.sin_port=htons((unsigned short)atoi(rem_service))) == 0) {
		throw SocketError("Can't get service entry", this);
	}

	rem_port = sin.sin_port;

    /* Map host name to IP address, allowing for dotted decimal */
	if (phe = gethostbyname(rem_host))
		memcpy(&sin.sin_addr, phe->h_addr, phe->h_length);
	else if ( (sin.sin_addr.s_addr = inet_addr(rem_host)) == INADDR_NONE ) {
		throw SocketError("Can't get host entry", this);
	}

    rem_ip = (char *) inet_ntoa(sin.sin_addr);

	/* Making a TCP socket */
	if ( (ppe = getprotobyname("tcp")) == 0)  {
		throw SocketError("Can't get protocol entry", this);
	}

	//type = SOCK_DGRAM;
	type = SOCK_STREAM;

    /* Allocate a socket */
	sck_id = socket(PF_INET, type, ppe->p_proto);
	if (sck_id < 0) {
		SocketError se("Unable to create socket", this);
		se.append_unix_error();
		throw se;
	}

    /* Connect the socket */
	if (connect(sck_id, (struct sockaddr *)&sin, sizeof(sin)) < 0) {
		SocketError se("Unable to connect to remote host", this);
		se.append_unix_error();
		throw se;
	}
}

////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////

ListenerSocket::ListenerSocket(const char *service, int qlen) {
	struct servent	*pse;	/* pointer to service information entry	*/
	struct protoent *ppe;	/* pointer to protocol information entry*/
	struct sockaddr_in sin;	/* an Internet endpoint address		*/
	int	type;				/* socket descriptor and socket type	*/

	int ret = 0;
	int value = 1;
	int vlen = sizeof(value);

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = INADDR_ANY;

    /* Map service name to port number */
	if ( pse = getservbyname(service, "tcp") )
		sin.sin_port = htons(ntohs((unsigned short)pse->s_port));
	else if ((sin.sin_port=htons((unsigned short)atoi(service))) == 0) {
		throw SocketError("Can't get service entry", this);
	}

    /* Map protocol name to protocol number */
	if ( (ppe = getprotobyname("tcp")) == 0) {
		throw SocketError("Can't get protocol entry", this);
	}

    /* Use protocol to choose a socket type */
	/*if (strcmp(transport, "udp") == 0)
		type = SOCK_DGRAM;
	else*/
	type = SOCK_STREAM;

    /* Allocate a socket */
	sck_id = socket(PF_INET, type, ppe->p_proto);
	if (sck_id < 0) {
		SocketError se("Can't create socket", this);
		se.append_unix_error();
		throw se;
	}

	// allow address reuse
	ret = setsockopt(sck_id, SOL_SOCKET, SO_REUSEADDR, &value, vlen);
	if(ret == -1) {
		SocketError se("Can't set socket options", this);
		se.append_unix_error();
		throw se;
	}


    /* Bind the socket */
	if (bind(sck_id, (struct sockaddr *)&sin, sizeof(sin)) < 0) {
		SocketError se("Can't bind to port", this);
		se.append_unix_error();
		throw se;
	}

	if (type == SOCK_STREAM && listen(sck_id, qlen) < 0) {
		SocketError se("Can't listen on port", this);
		se.append_unix_error();
		throw se;
	}
}

ServerSocket *ListenerSocket::accept_connection() {
    struct sockaddr_in addr;

    socklen_t flen;
    int newsck;

    flen = sizeof addr;
    newsck = accept(sck_id, (struct sockaddr*) &addr, &flen);
    if (newsck <= 0) {
        SocketError se("Error accepting connection", this);
        se.append_unix_error();
        throw se;
    }

    ServerSocket *ss = new ServerSocket(newsck);
    char buf[16];

    // Retrieve remote address
    memset((void *) buf, 0, 16);
    strcpy(buf, (char *) inet_ntoa(addr.sin_addr));
    ss->rem_ip = buf;

    // Retrieve remote port
    ss->rem_port = (unsigned int) ntohs(addr.sin_port);

    // Retireve local IP address
	flen = sizeof addr;
	getsockname(newsck, (struct sockaddr*) &addr, &flen);

	memset((void *) buf, 0, 16);
	strcpy(buf, (char *) inet_ntoa(addr.sin_addr));
	ss->local_ip = buf;

    return ss;
}

pid_t ListenerSocket::accept_connection_and_fork(ServerSocket * &ss) {
    struct sockaddr_in addr;

    socklen_t flen;
    int newsck;

    flen = sizeof addr;
    newsck = accept(sck_id, (struct sockaddr*) &addr, &flen);
    if (newsck <= 0) {
        SocketError se("Error accepting connection", this);
        se.append_unix_error();
        throw se;
    }

    pid_t f_ret = fork();

    // if not child...
    if (f_ret != 0) {
    	ss = NULL;
    	close(newsck);
    	return f_ret;
    }

    ss = new ServerSocket(newsck);

    char buf[16];

    // Retrieve remote address
    memset((void *) buf, 0, 16);
    strcpy(buf, (char *) inet_ntoa(addr.sin_addr));
    ss->rem_ip = buf;

    // Retrieve remote port
    ss->rem_port = (unsigned int) ntohs(addr.sin_port);

    // Retireve local IP address
	flen = sizeof addr;
	getsockname(newsck, (struct sockaddr*) &addr, &flen);

	memset((void *) buf, 0, 16);
	strcpy(buf, (char *) inet_ntoa(addr.sin_addr));
	ss->local_ip = buf;

    return f_ret;
}
